package edu.indiana.d2i.vmm.cluster;

import java.io.BufferedInputStream;
import java.io.BufferedOutputStream;
import java.io.ByteArrayInputStream;
import java.io.FileOutputStream;
import java.io.FileWriter;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.io.UnsupportedEncodingException;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;

import com.jcraft.jsch.Channel;
import com.jcraft.jsch.ChannelExec;
import com.jcraft.jsch.ChannelSftp;
import com.jcraft.jsch.JSch;
import com.jcraft.jsch.JSchException;
import com.jcraft.jsch.Session;
import com.jcraft.jsch.SftpException;
import com.jcraft.jsch.SftpProgressMonitor;

public class SSHSession {
	private JSch jsch = new JSch();
	private Session session;
	private StringBuilder cmdPool = new StringBuilder();
	private String host;
	private final int MAX_TRY = 5;

	// private final String SUDO_PREFIX = "sudo -S -p '' ";
	private final String SUDO_PREFIX = "sudo ";

	private static final Log logger = LogFactory.getLog(SSHSession.class);

	private int getCmdChannelScreenOutput(ChannelExec channel)
			throws IOException {
		InputStream in = channel.getInputStream();
		// don't disconnect too early!
		byte[] tmp = new byte[1024];
		while (true) {
			while (in.available() > 0) {
				int i = in.read(tmp, 0, 1024);
				if (i < 0)
					break;
				System.out.print(new String(tmp, 0, i));
			}
			if (channel.isClosed()) {
				System.out.println("exit-status: " + channel.getExitStatus());
				return channel.getExitStatus();
			}
			try {
				Thread.sleep(1000);
			} catch (Exception ee) {
			}
		}
	}

	public SSHSession(String host, String user, String passwd)
			throws JSchException {
		this.host = host;
		java.util.Properties config = new java.util.Properties();
		config.put("StrictHostKeyChecking", "no");
		session = jsch.getSession(user, host, 22);
		session.setPassword(passwd);
		session.setConfig(config);
		session.connect();
	}

	public void close() throws Exception {
		flush();
		session.disconnect();
	}

	public String getHost() {
		return host;
	}

	public int flush() throws Exception {
		int exitcode = 0;
		if (cmdPool.length() > 0) {
			Channel channel = session.openChannel("exec");
			((ChannelExec) channel).setCommand(cmdPool.toString());
			((ChannelExec) channel).setErrStream(System.err);

			channel.connect();
			exitcode = getCmdChannelScreenOutput((ChannelExec) channel);
			channel.disconnect();
		}
		cmdPool.delete(0, cmdPool.length());
		return exitcode;
	}

	public void writeSmallFile(String content, String pathname, boolean sudo)
			throws Exception {
		writeFile(new ByteArrayInputStream(content.getBytes("UTF-8")),
				pathname, sudo);
	}

	public void writeFile(InputStream input, String pathname, boolean sudo)
			throws Exception {
		ChannelSftp sftpchannel = (ChannelSftp) session.openChannel("sftp");
		sftpchannel.connect();
		if (sudo) {
			Channel cmdchannel = session.openChannel("exec");
			((ChannelExec) cmdchannel).setCommand(SUDO_PREFIX + "rm -fr "
					+ pathname + ";" + SUDO_PREFIX + "touch " + pathname + ";"
					+ SUDO_PREFIX + "chmod o+w " + pathname + ";" + "exit;");
			cmdchannel.connect();
			getCmdChannelScreenOutput((ChannelExec) cmdchannel);
			cmdchannel.disconnect();
		}
		sftpchannel.put(input, pathname, ChannelSftp.OVERWRITE);
		sftpchannel.disconnect();
	}

	public void execCmd(String cmd, boolean sudo) {
		if (cmdPool.length() > 0)
			cmdPool.append(";");
		if (sudo)
			cmdPool.append(SUDO_PREFIX + cmd);
		else
			cmdPool.append(cmd);
	}

	public void execCmdWithoutBlocking(String cmd, boolean sudo)
			throws Exception {
		String command = "";
		if (sudo)
			command = SUDO_PREFIX + cmd;
		else
			command = cmd;
		Channel channel = session.openChannel("exec");
		((ChannelExec) channel).setCommand(command);
		Thread.sleep(2000);

		channel.connect();
		channel.disconnect();
	}

	public StringBuffer execCmdReturnScreenOutput(String cmd, boolean sudo)
			throws Exception {
		StringBuffer buf = new StringBuffer();

		Channel channel = session.openChannel("exec");
		((ChannelExec) channel).setCommand(cmd);
		((ChannelExec) channel).setErrStream(System.err);

		channel.connect();
		InputStream in = channel.getInputStream();
		// don't disconnect too early!
		byte[] tmp = new byte[1024];
		while (true) {
			while (in.available() > 0) {
				int i = in.read(tmp, 0, 1024);
				if (i < 0)
					break;
				buf.append(new String(tmp, 0, i));
			}
			if (channel.isClosed()) {
				System.out.println("exit-status: " + channel.getExitStatus());
				break;
			}
			Thread.sleep(1000);
		}
		channel.disconnect();

		return buf;
	}

	public void downloadFile(String rfile, String lfile) throws Exception {
		ChannelSftp sftpchannel = (ChannelSftp) session.openChannel("sftp");
		sftpchannel.connect();

		BufferedOutputStream output = new BufferedOutputStream(
				new FileOutputStream(lfile));

		BufferedInputStream inputStream = new BufferedInputStream(
				sftpchannel.get(rfile));
		byte[] buf = new byte[2048];
		int length = 0;
		while ((length = inputStream.read(buf)) != -1) {
			output.write(buf, 0, length);
		}
		inputStream.close();
		output.close();
		sftpchannel.disconnect();
	}
}
